import torch
import cv2
from interface import MolScribe
# from huggingface_hub import hf_hub_download
 
# 下载 MolScribe 的预训练模型
# ckpt_path = hf_hub_download('yujieq/MolScribe', 'swin_base_char_aux_1m.pth')
ckpt_path = 'output/swin_base_lstm_best.pth'

# 初始化模型
model = MolScribe(ckpt_path, device=torch.device('cpu'))

 
# 预测分子结构
output = model.predict_image_file('assets/e.png', return_atoms_bonds=True, return_confidence=True)
 
# 输出结果
print(output)  

# image = cv2.imread('images.png')
# result_image = model.draw_prediction(output, image)
# cv2.imshow("Prediction", result_image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()






